from matplotlib.axes import Axes
import numpy as np
from environment.ElectricScheduleEnv import env
import einops
from environment.ElectricScheduleEnv import ElectricScheduleEnv
import torch.nn as nn


def hard_update(target, source):
    """
    Copy network parameters from source to target
    Inputs:
        target (torch.nn.Module): Net to copy parameters to
        source (torch.nn.Module): Net whose parameters to copy
    """
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(param.data)


def split_action(actions: tuple, actions_dims_dict: dict, ts: int):
    i = 0
    structured_action = []
    storage = []
    area = -1
    for action in actions:
        area += 1
        action = Scaler.transform(action, i)
        # fire_dim = actions_dims_dict['fire'][i]
        fire_dim = 1
        storage_dim = actions_dims_dict['storage'][i]
        dim = actions_dims_dict['total'][i]
        a = {
            # "fire": action[:fire_dim],
            "fire": action[0],
            "storage": action[fire_dim:storage_dim+fire_dim],
            # "outer": action[storage_dim+fire_dim],
            # "exchange": action[storage_dim+fire_dim+1],
            "exchange": 0,
            "buy": 0,
            "solar": 0,
            "wind": 0,
        }
        tot = action[0]
        # add new energy output
        if area == 0:
            tot += env.A.Wind.PredictOutput[ts] + env.A.Solar.PredictOutput[ts]
        if area == 1:
            tot += env.B.Wind.PredictOutput[ts] + env.B.Solar.PredictOutput[ts]
        for s in a['storage']:
            tot += s
        storage.append(tot)

        structured_action.append(a)
        i += 1

    structured_action[0]['solar'] = env.A.Solar.PredictOutput[ts]
    structured_action[0]['wind'] = env.A.Wind.PredictOutput[ts]
    structured_action[1]['solar'] = env.B.Solar.PredictOutput[ts]
    structured_action[1]['wind'] = env.B.Wind.PredictOutput[ts]

    loadA, loadB = env.A.Load[ts], env.B.Load[ts]
    totalA, totalB = storage
    if totalA >= loadA and totalB >= loadB:
        fireA, fireB = structured_action[0]["fire"], structured_action[1]["fire"]
        extraA, extraB = totalA - loadA, totalB - loadB
        # structured_action[0]["sell"] += extraA
        # structured_action[1]["sell"] += extraB
        structured_action[0]["fire"] -= (extraA if fireA >= extraA else fireA)
        structured_action[1]["fire"] -= (extraB if fireB >= extraB else fireB)
        print(f"time at {ts} step here 1")
        if structured_action[0]["fire"] == 0:
            extraA -= fireA
            structured_action[0]["exchange"] += extraA
            structured_action[1]["exchange"] -= extraA
            structured_action[1]["fire"] -= extraA
            print(f"time at {ts} step here 2")
        if structured_action[1]["fire"] == 0:
            extraB -= fireB
            structured_action[1]["exchange"] += extraB
            structured_action[0]["exchange"] -= extraB
            structured_action[0]["fire"] -= extraB
            print(f"time at {ts} step here 3")
        return structured_action
    if totalA >= loadA and totalB <= loadB:
        leftA = totalA - loadA
        requireB = loadB - totalB
        if leftA >= requireB:
            structured_action[0]["exchange"] = requireB
            structured_action[1]["exchange"] = -requireB
            structured_action[0]["fire"] -= leftA-requireB
            print(f"time at {ts} step here 4")
            # structured_action[0]["storage"] -= leftA-requireB
        else:
            structured_action[0]["exchange"] += leftA
            structured_action[1]["exchange"] -= leftA
            requireB -= leftA
            canA = env.A.FireMaxOutput - structured_action[0]["fire"]
            if canA >= requireB:
                structured_action[0]["fire"] += requireB
                structured_action[0]["exchange"] += requireB
                structured_action[1]["exchange"] -= requireB
                print(f"time at {ts} step here 5")
            else:
                structured_action[0]["fire"] += canA
                structured_action[0]["exchange"] += canA
                structured_action[1]["exchange"] -= canA
                structured_action[1]["fire"] += requireB - canA
                print(f"time at {ts} step here 6")
        return structured_action
    if totalA <= loadA and totalB >= loadB:
        leftB = totalB - loadB
        requireA = loadA - totalA
        if leftB >= requireA:
            structured_action[0]["exchange"] = -requireA
            structured_action[1]["exchange"] = requireA
            structured_action[1]["fire"] -= leftB-requireA
            print(f"time at {ts} step here 7")
        else:
            structured_action[0]["exchange"] -= leftB
            structured_action[1]["exchange"] += leftB
            requireA -= leftB
            canB = env.B.FireMaxOutput - structured_action[1]["fire"]
            if canB >= requireA:
                structured_action[1]["fire"] += requireA
                structured_action[1]["exchange"] += requireA
                structured_action[0]["exchange"] -= requireA
                print(f"time at {ts} step here 8")
            else:
                structured_action[0]["fire"] += canA
                structured_action[0]["exchange"] += canA
                structured_action[1]["exchange"] -= canA
                structured_action[1]["fire"] += requireB - canA
                print(f"time at {ts} step here 9")
        return structured_action

    if totalA <= loadA and totalB <= loadB:
        requireA = loadA - totalA
        requireB = loadB - totalB
        fireA = structured_action[0]["fire"]
        fireB = structured_action[1]['fire']
        fireLeftA = env.A.FireMaxOutput - fireA
        fireLeftB = env.B.FireMaxOutput - fireB
        # 本情景只需要考虑A不足
        if fireLeftA >= requireA:
            structured_action[0]['fire'] += requireA
            print(f"time at {ts} step here 10")
        else:
            structured_action[0]['fire'] += fireLeftA
            structured_action[1]['fire'] += requireA-fireLeftA+requireB
            structured_action[1]["exchange"] += requireA-fireLeftA
            structured_action[0]["exchange"] -= requireA-fireLeftA
            print(f"time at {ts} step here 11")
        return structured_action

    # print(f"time={ts}\tfire={structured_action[0]['fire']}\twind={env.A.Wind.PredictOutput[ts]}\tsolar={env.A.Solar.PredictOutput[ts]}\texchage={}")


def print_net(net):
    for i in range(len(net.agents)):
        print("===============agent-%d================" % (i+1))
        print(net.agents[i].policy)
        print("===============critic-%d================" % (i+1))
        print(net.agents[i].critic)


def env_info(env: ElectricScheduleEnv):
    print("+-----------Day Ahead Power System Schedule-----------+")
    print("| State Dim =", env.state_dim, "                                 |")
    print("| Action Dim =", env.action_dim,
          "                                |")
    print("+-----------------------------------------------------+")


def disable_gradients(module: nn.Module):
    for p in module.parameters():
        p.requires_grad = False


def enable_gradients(module):
    for p in module.parameters():
        p.requires_grad = True


# 将[-1, 1]的动作空间映射到指定的范围
class ActionScaler:
    def __init__(self, env: ElectricScheduleEnv) -> None:
        self.env = env

    def transform(self, action, agent_id: int):
        action_dim = self.env.action_dim[agent_id]
        origin_shape = action.shape
        action = action.reshape(-1, action_dim)
        action = einops.rearrange(action, 'b a -> a b')
        area = self.env.Areas[agent_id]
        index = 0
        # '''area.FireMinOutput +'''
        action[index] = (action[index] + 1) * (area.FireMaxOutput) * 0.5
        index += 1
        # for fire in range(index, index + area.FireStationsSize):
        #     fire_min = area.FireStations[fire-index].OutputRange[0]
        #     fire_max = area.FireStations[fire-index].OutputRange[3]
        #     action[fire] = fire_min + \
        #         (action[fire] + 1) * (fire_max - fire_min) * 0.5
        # index += area.FireStationsSize
        for storage in range(index, index + area.StorageStationsSize):
            k = area.StorageStations[storage-index].Output
            action[storage] *= k
        index += area.StorageStationsSize
        # action[index] = self.env.Link.AOutPower * action[index]
        # index += 1
        # action[index] = self.env.Link.InternalExchangePower * action[index]
        # index += 1
        # solar scale
        # action[index] = (action[index] + 1) * \
        #     area.Solar.InstalledCapacity * 0.5
        # # wind scale
        # action[index+1] = (action[index+1] + 1) * \
        #     area.Wind.InstalledCapacity * 0.5
        # invert transform
        action = einops.rearrange(action, 'a b -> b a')
        action = action.reshape(origin_shape)
        return action


def draw_policy(ax: Axes, index: int, actions: dict, area: int, exchange: float):
    _bottom, bottom = 0, 0
    storage = np.sum(actions["storage"])
    if storage > 0:
        # 放电
        ax.bar(index, storage, color="violet", label="储能")
        bottom += storage
    else:
        ax.bar(index, storage, color="violet", label="储能")
        _bottom += storage

    if exchange > 0:  # 向外输电
        ax.bar(index, -exchange, bottom=_bottom, color="orange", label="联络线")
        _bottom -= exchange
    else:
        ax.bar(index, -exchange, bottom=bottom, color="orange", label="联络线")
        bottom -= exchange

    buy = actions["buy"]
    ax.bar(index, buy, bottom=bottom, color="lightpink", label="购电")
    bottom += actions["buy"]

    solar = env.Areas[area].Solar.PredictOutput[index]
    wind = env.Areas[area].Wind.PredictOutput[index]

    actionA_fire_sum = np.sum(actions['fire'])
    ax.bar(index, actionA_fire_sum, bottom=bottom,
           color="goldenrod", label="火电")
    bottom += actionA_fire_sum

    ax.bar(index, solar, bottom=bottom, color="lightseagreen", label="太阳能")
    bottom += solar

    ax.bar(index, wind, bottom=bottom, color="cornflowerblue", label="风能")


Scaler = ActionScaler(env)


def calc_output(action: dict):
    tot = action['fire'] + action['wind'] + action['solar']
    for s in action['storage']:
        tot -= s
    tot -= action['exchange']
    return tot
